module Zeno.FunctionDef where

import Prelude ()
import StdImports
import qualified Data.Set as Set
import Zeno.Term

data FunctionDef a 
  = FunctionDef { funcName :: !a,
                  funcParameters :: ![a],
                  funcExpression :: !(Expression a) }
  deriving ( Eq )
                
data Expression a 
  = CaseOf !(Term a) ![When a]
  | Value !(Term a)
  deriving ( Eq )
  
data When a
  = When !(Term a) !(Expression a)
  deriving ( Eq )
  
instance NFData a => NFData (FunctionDef a) where
  rnf (FunctionDef n ps ex) = rnf n `seq` rnf ps `seq` rnf ex
  
instance NFData a => NFData (Expression a) where
  rnf (Value t) = rnf t
  rnf (CaseOf t ws) = rnf t `seq` rnf ws
  
instance NFData a => NFData (When a) where
  rnf (When t ex) = rnf t `seq` rnf ex
  
instance Functor FunctionDef where
  fmap f (FunctionDef n ps ex) = 
    FunctionDef (f n) (map f ps) (fmap f ex)
    
instance Foldable FunctionDef where
  foldMap f (FunctionDef n ps ex) = 
    f n `mappend` foldMap f ps `mappend` foldMap f ex
    
instance Traversable FunctionDef where
  sequenceA (FunctionDef n ps ex) = 
    FunctionDef <$> n <*> sequenceA ps <*> sequenceA ex
  
instance Functor Expression where
  fmap f (CaseOf t ws) = 
    CaseOf (fmap f t) (map (fmap f) ws)
  fmap f (Value t) = Value (fmap f t)
  
instance Foldable Expression where
  foldMap f (CaseOf t ws) =
    foldMap f t `mappend` foldMap (foldMap f) ws
  foldMap f (Value t) = 
    foldMap f t
    
instance Traversable Expression where
  sequenceA (CaseOf t ws) =
    CaseOf <$> sequenceA t <*> sequenceA (map sequenceA ws)
  sequenceA (Value t) =
    Value <$> sequenceA t
  
instance Functor When where
  fmap f (When t expr) = 
    When (fmap f t) (fmap f expr)
   
instance Foldable When where
  foldMap f (When t expr) =
    foldMap f t `mappend` foldMap f expr
    
instance Traversable When where
  sequenceA (When t expr) = 
    When <$> sequenceA t <*> sequenceA expr
    
instance TermTraversable Expression where
  mapTermsM f (CaseOf t ws) = do
    ws' <- mapM (mapTermsM f) ws
    t' <- mapTermsM f t
    return (CaseOf t' ws')
  mapTermsM f (Value t) =
    Value `liftM` (mapTermsM f t)
    
instance TermTraversable When where
  mapTermsM f (When t expr) =
    When `liftM` (mapTermsM f t) `ap` (mapTermsM f expr)
    
instance Show a => Show (FunctionDef a) where
  show (FunctionDef name params expr) =
    let params' = intercalate " " (map show params)
        expr' = runReader (showExpr expr) 1
        name' = show name
    in name' ++ " " ++ params' ++ " = " ++ expr'
   
    
showWhen :: Show a => When a -> Indented String
showWhen (When term expr) = do
  i <- indentation
  expr' <- indent $ showExpr expr
  return $ i ++ (show term) ++ " -> " ++ expr' 

instance Show a => Show (Expression a) where
  show = (flip runReader 0) . showExpr 
  
showExpr :: Show a => Expression a -> Indented String
showExpr (Value t) = do
  i <- indentation
  return $ i ++ show t
showExpr (CaseOf t ws) = do
  i <- indentation
  ws' <- indent $ mapM showWhen ws
  return $ i ++ "case " ++ (show t) ++ " of" ++ (concat ws')


